from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, IntegerType, StringType

if __name__ == '__main__':
    # 构建SparkSession对象
    spark = SparkSession.builder. \
        appName("local[*]"). \
        config("spark.sql.shuffle.partitions", "4"). \
        getOrCreate()
    # appName 设置程序名称
    # config 设置常用属性。可以通过此来设置配置
    # 最后通过getOrCreate 创建 SparkSession对象

    # 从SparkSession中获取SparkContext
    sc = spark.sparkContext

    # 首先构建一个RDD
    rdd = sc.textFile("../../data/sql/stu_score.txt"). \
        map(lambda line: line.split(",")). \
        map(lambda x: (int(x[0]), x[1], int(x[2])))  # 需要做类型转换，因为类型从RDD中推断，默认为string

    # 通过StructType来定义表结构
    schema = StructType(). \
        add("id", IntegerType(), nullable=False). \
        add("name", StringType(), nullable=True). \
        add("score", IntegerType(), nullable=False)
    # 一个add定义一个列的信息
    # add方法：参数1：列名称；参数2：列类型；参数3：是否允许为空

    # 通过rdd.toDF()方法获取DataFrame对象
    # 方式1：只传列名，类型靠推断
    df = rdd.toDF(['id', 'name', 'score'])
    df.printSchema()
    df.show(truncate=False)

    # 方式1：传入StructType对象
    df = rdd.toDF(schema)
    df.printSchema()
    df.show(truncate=False)

    # df.createTempView("score")
    #
    # spark.sql("select * from score where score < 99").show()
